Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fixes Imagen sampling example and updates container #868

Merged
merged 13 commits into from
Aug 26, 2024
Merged

Fixes Imagen sampling example and updates container #868

merged 13 commits into from
Aug 26, 2024

Conversation

terrykong
Copy link
Contributor

@terrykong terrykong commented May 30, 2024 •
edited
Loading

  1. Fix imagen sampling loop when prompt_ct is a multiple of batch_size // gen_per_prompt
  2. Add comment explaining the invocation of sampling script of what exactly is expected due to implicit checkpoint dir requirements and quoting
  3. add 2B base model generation gin configs
  4. parametrize imagen sampling scripts
  5. Updates the imagen image with the fix in (1); built as follows:
docker buildx build --push -t ghcr.io/nvidia/t5x:imagen-2023-10-02.v3 ./JAX-Toolbox -f - <<EOF
FROM ghcr.io/nvidia/t5x:imagen-2023-10-02
COPY rosetta/rosetta/projects/diffusion/common/set_gpu_xla_flags.sh /opt/rosetta/rosetta/projects/diffusion/common/set_gpu_xla_flags.sh
COPY rosetta/rosetta/projects/imagen/README.md /opt/rosetta/rosetta/projects/imagen/README.md
COPY rosetta/rosetta/projects/imagen/configs/imagen_1024_sample_2b.gin /opt/rosetta/rosetta/projects/imagen/configs/imagen_1024_sample_2b.gin
COPY rosetta/rosetta/projects/imagen/configs/imagen_256_sample_2b.gin /opt/rosetta/rosetta/projects/imagen/configs/imagen_256_sample_2b.gin
COPY rosetta/rosetta/projects/imagen/imagen_pipe.py /opt/rosetta/rosetta/projects/imagen/imagen_pipe.py
COPY rosetta/rosetta/projects/imagen/scripts/example_slurm_inf_train.sub /opt/rosetta/rosetta/projects/imagen/scripts/example_slurm_inf_train.sub
COPY rosetta/rosetta/projects/imagen/scripts/sample_imagen_1024.sh /opt/rosetta/rosetta/projects/imagen/scripts/sample_imagen_1024.sh
COPY rosetta/rosetta/projects/imagen/scripts/sample_imagen_256.sh /opt/rosetta/rosetta/projects/imagen/scripts/sample_imagen_256.sh
EOF

@terrykong terrykong marked this pull request as draft June 5, 2024 23:48
@terrykong
Fix imagen loop when prompt_ct is a multiple of
cc4290b 
batch_size//gen_per_prompt
@terrykong
Add a comment explaining the invocation of the sampling script for
c35a14f 
imagen
@terrykong
Update imagen image with this fix:
c0d4885 
```
docker buildx build --push -t ghcr.io/nvidia/t5x:imagen-2023-10-02.v2 . -f - <<EOF
FROM ghcr.io/nvidia/t5x:imagen-2023-10-02
COPY rosetta/rosetta/projects/imagen/imagen_pipe.py /opt/rosetta/rosetta/projects/imagen/imagen_pipe.py
EOF
```
@terrykong
parametrize batch size for imagen
4caed81
@terrykong
update documentation with new keys and values
86511b2
@terrykong
add a timer
7c55439 
docker buildx build --push -t ghcr.io/nvidia/t5x:imagen-2023-10-02.v2 .
-f - <<EOF
FROM ghcr.io/nvidia/t5x:imagen-2023-10-02
COPY rosetta/rosetta/projects/imagen/README.md
/opt/rosetta/rosetta/projects/imagen/README.md
COPY rosetta/rosetta/projects/imagen/imagen_pipe.py
/opt/rosetta/rosetta/projects/imagen/imagen_pipe.py
COPY rosetta/rosetta/projects/imagen/scripts/sample_imagen_1024.sh
/opt/rosetta/rosetta/projects/imagen/scripts/sample_imagen_1024.sh
COPY rosetta/rosetta/projects/imagen/scripts/sample_imagen_256.sh
/opt/rosetta/rosetta/projects/imagen/scripts/sample_imagen_256.sh
EOF
@terrykong
typo
0c235e3
@terrykong
better logging
3cf8c94
@terrykong
update readme on the divisible by gpu and add MAX_FILENAME_LENGTH
cc753e1
@terrykong
typo in perf calc
9940b9c
@terrykong
reset time
6baba82
@terrykong
few other missing imagen tags updated
c17c496 
Signed-off-by: Terry Kong <terryk@nvidia.com>
@terrykong terrykong marked this pull request as ready for review July 22, 2024 23:13
@terrykong terrykong merged commit de12e17 into main Aug 26, 2024
151 of 156 checks passed
@terrykong terrykong deleted the imagen-fix branch August 26, 2024 21:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ashors1 ashors1 ashors1 approved these changes

@SahilJain314 SahilJain314 Awaiting requested review from SahilJain314

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@terrykong @ashors1